#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <string.h>
#include <assert.h>
#include <mpi.h>
#include <fitsio.h>

#include "halofinder.h"

static char *ttype[] = { "haloid", "radius", "mass", "vcirc", "px", "py", "pz", "vx", "vy", "vz"};
static char *tform[] = { "K", "E", "E", "E", "E", "E", "E", "E", "E", "E"};

static int appendpeakshalos2file(long firstRow, long NumBuffHalos, Halo *buffHalos, char *buff, fitsfile *fptr, int *status);

int writepeakshalos_name(char *name)
{
  fitsfile *fptr;
  int tfields,status,ext;
  char *buff;
  LONGLONG nrows;
  long chunkInd,firstInd,lastInd,NumHalosInChunkBase,NumHalosInChunk,NumChunks;
  int currTask;
  Halo *buffHalos;
  long firstRow,i;
  char bangName[MAX_FILENAME];
  
  status = 0;
  ext = 2;
  
  //make the file
  if(ThisTask == 0)
    {
      sprintf(bangName,"!%s",name);
      
      fits_create_file(&fptr,bangName,&status);
      if(status)
	fits_report_error(stderr,status);
      
      nrows = 0;
      tfields = 10;
      fits_create_tbl(fptr,BINARY_TBL,nrows,tfields,ttype,tform,NULL,NULL,&status);
      if(status)
	fits_report_error(stderr,status);
      
      fits_close_file(fptr,&status);
      if(status)
	fits_report_error(stderr,status);

#ifdef DEBUG
      fprintf(stderr,"%d: initialized fits file\n",ThisTask);
#endif
    }
  
  for(currTask=0;currTask<NTasks;++currTask)
    {
      MPI_Barrier(MPI_COMM_WORLD);
      
      if(ThisTask == currTask)
	{

#ifdef DEBUG
	  fprintf(stderr,"%d: writing %ld halos to fits file\n",ThisTask,NumHalos);
#endif
	  
	  if(NumHalos > 0)
	    {
	      //open file and get number of rows
	      fits_open_table(&fptr,name,READWRITE,&status);
	      if(status)
		fits_report_error(stderr,status);
	      
	      fits_movabs_hdu(fptr,ext,NULL,&status);
	      if(status)
		fits_report_error(stderr,status);
	      
	      //NumHalosInChunkBase = (long) ((FITSIO_CHUNK_SIZE_MB*1048576.0)/(sizeof(Halo)));
	      fits_get_rowsize(fptr,&NumHalosInChunkBase,&status);
	      if(status)
		fits_report_error(stderr,status);
	      NumChunks = NumHalos/NumHalosInChunkBase;
	      if(NumChunks*NumHalosInChunkBase < NumHalos)
		NumChunks += 1;
	      
#ifdef DEBUG
	      if(DEBUG_LEVEL > 0)
		fprintf(stderr,"%d: NumHalos = %ld, NumHalosInChunkBase = %ld, NumChunks = %ld\n",ThisTask,
			NumHalos,NumHalosInChunkBase,NumChunks);
#endif
	      
	      if(sizeof(long) > sizeof(float))
		buff = (char*)malloc(sizeof(long)*NumHalosInChunkBase);
	      else
		buff = (char*)malloc(sizeof(float)*NumHalosInChunkBase);
	      assert(buff != NULL);
	      
	      buffHalos = (Halo*)malloc(sizeof(Halo)*NumHalosInChunkBase);
	      assert(buffHalos != NULL);
	      
	      fits_get_num_rowsll(fptr,&nrows,&status);
	      if(status)
		fits_report_error(stderr,status);
	      
	      firstRow = (long) nrows;
	      
	      for(chunkInd=0;chunkInd<NumChunks;++chunkInd)
		{
		  firstInd = chunkInd*NumHalosInChunkBase;
		  lastInd = (chunkInd+1)*NumHalosInChunkBase-1;
		  if(lastInd >= NumHalos-1)
		    lastInd = NumHalos-1;
		  NumHalosInChunk = lastInd - firstInd + 1;
		  
		  for(i=0;i<NumHalosInChunk;++i)
		    buffHalos[i] = Halos[i+firstInd];
		  
#ifdef DEBUG
		  if(DEBUG_LEVEL > 1)
		    fprintf(stderr,"%d: chunkInd = %ld, firstInd,lastInd,NumHalosInChunk = %ld|%ld|%ld\n",ThisTask,
			    chunkInd,firstInd,lastInd,NumHalosInChunk);
#endif
		  
		  appendpeakshalos2file(firstRow,NumHalosInChunk,buffHalos,buff,fptr,&status);
		  firstRow += NumHalosInChunk;
		  if(status)
		    fits_report_error(stderr,status);
		}
	      
	      //close file
	      fits_close_file(fptr,&status);
	      if(status)
		fits_report_error(stderr,status);
	      
	      free(buff);
	      free(buffHalos);
	    }
	}
      
      MPI_Barrier(MPI_COMM_WORLD);
    }
  
  return status;
}

static int appendpeakshalos2file(long firstRow, long NumBuffHalos, Halo *buffHalos, char *buff, fitsfile *fptr, int *status)
{
  int colnum;
  long *larr,i,j;
  float *farr;
  LONGLONG firstrow,firstelem,nelements;
  
  firstrow = (LONGLONG) (firstRow+1);
  firstelem = 1;
  nelements = (LONGLONG) (NumBuffHalos);
  
  colnum = 1;
  larr = (long*)buff;
  for(j=0;j<NumBuffHalos;++j)
    larr[j] = buffHalos[j].index;
  fits_write_col(fptr,TLONG,colnum,firstrow,firstelem,nelements,larr,status);
  if(*status)
    fits_report_error(stderr,*status);
  ++colnum;
  
  farr = (float*)buff;
  
  for(j=0;j<NumBuffHalos;++j)
    farr[j] = buffHalos[j].radius;
  fits_write_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,farr,status);
  if(*status)
    fits_report_error(stderr,*status);
  ++colnum;
  
  for(j=0;j<NumBuffHalos;++j)
    farr[j] = buffHalos[j].mass;
  fits_write_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,farr,status);
  if(*status)
    fits_report_error(stderr,*status);
  ++colnum;
  
  for(j=0;j<NumBuffHalos;++j)
    farr[j] = buffHalos[j].vcirc;
  fits_write_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,farr,status);
  if(*status)
    fits_report_error(stderr,*status);
  ++colnum;
  
  for(i=0;i<3;++i)
    {
      for(j=0;j<NumBuffHalos;++j)
	farr[j] = buffHalos[j].pos[i];
      fits_write_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,farr,status);
      if(*status)
	fits_report_error(stderr,*status);
      ++colnum;
    }
      
  for(i=0;i<3;++i)
    {
      for(j=0;j<NumBuffHalos;++j)
	farr[j] = buffHalos[j].vel[i];
      fits_write_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,farr,status);
      if(*status)
	fits_report_error(stderr,*status);
      ++colnum;
    }
  
  return *status;
}

long read_totnumpeakshalos_fits(char name[])
{
  fitsfile *fptr;
  int status,ext;
  long TotNumHalos;
  
  status = 0;
  ext = 2;
  
  fits_open_table(&fptr,name,READONLY,&status);
  if(status)
    fits_report_error(stderr,status);
      
  fits_movabs_hdu(fptr,ext,NULL,&status);
  if(status)
    fits_report_error(stderr,status);
  
  fits_get_num_rows(fptr,&TotNumHalos,&status);
  if(status)
    fits_report_error(stderr,status);
  
  fits_close_file(fptr,&status);
  if(status)
    fits_report_error(stderr,status);
  
  
  return TotNumHalos;
}

long read_chunksizepeakshalos_fits(char name[])
{
  fitsfile *fptr;
  int status,ext;
  long NumHalosInChunkBase;
  
  status = 0;
  ext = 2;
  
  fits_open_table(&fptr,name,READONLY,&status);
  if(status)
    fits_report_error(stderr,status);
      
  fits_movabs_hdu(fptr,ext,NULL,&status);
  if(status)
    fits_report_error(stderr,status);
  
  fits_get_rowsize(fptr,&NumHalosInChunkBase,&status);
  if(status)
    fits_report_error(stderr,status);
  
  fits_close_file(fptr,&status);
  if(status)
    fits_report_error(stderr,status);
  
  return NumHalosInChunkBase;
}

void read_peakshaloschunk_fits(char name[], Halo *chunkHalos, long BaseHaloChunkSize, long HaloChunkSize, long chunkInd)
{
  fitsfile *fptr;
  int colnum,colnum_ttype,status,ext,anynul;
  long *larr,i,j;
  float *farr,nulval;
  LONGLONG firstrow,firstelem,nelements;
  long firstInd;
  char *buff;
  
  status = 0;
    
  nulval = 0.0;
  ext = 2;
  firstelem = 1;
      
  fits_open_table(&fptr,name,READONLY,&status);
  if(status)
    fits_report_error(stderr,status);
  
  fits_movabs_hdu(fptr,ext,NULL,&status);
  if(status)
    fits_report_error(stderr,status);
  
  if(sizeof(long) > sizeof(float))
    buff = (char*)malloc(sizeof(long)*BaseHaloChunkSize);
  else
    buff = (char*)malloc(sizeof(float)*BaseHaloChunkSize);
  assert(buff != NULL);
  farr = (float*)buff;
  larr = (long*)buff;
      
  firstInd = chunkInd*BaseHaloChunkSize;
  firstrow = (LONGLONG) (firstInd+1);
  nelements = (LONGLONG) (HaloChunkSize);
  
  //read index
  colnum_ttype = 0;
  fits_get_colnum(fptr,CASEINSEN,ttype[colnum_ttype],&colnum,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' column # = %d\n",ttype[colnum_ttype],colnum);
#endif
  fits_read_col(fptr,TLONG,colnum,firstrow,firstelem,nelements,&nulval,larr,&anynul,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' anynul = %d\n",ttype[colnum_ttype],anynul);
#endif
  for(j=0;j<HaloChunkSize;++j)
    chunkHalos[j].index = larr[j];
  ++colnum_ttype;
	  
  //read radius
  fits_get_colnum(fptr,CASEINSEN,ttype[colnum_ttype],&colnum,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' column # = %d\n",ttype[colnum_ttype],colnum);
#endif
  fits_read_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,&nulval,farr,&anynul,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' anynul = %d\n",ttype[colnum_ttype],anynul);
#endif
  for(j=0;j<HaloChunkSize;++j)
    chunkHalos[j].radius = farr[j];
  ++colnum_ttype;

  //read peakOverdens or mass
  fits_get_colnum(fptr,CASEINSEN,ttype[colnum_ttype],&colnum,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' column # = %d\n",ttype[colnum_ttype],colnum);
#endif
  fits_read_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,&nulval,farr,&anynul,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' anynul = %d\n",ttype[colnum_ttype],anynul);
#endif
  for(j=0;j<HaloChunkSize;++j)
    chunkHalos[j].mass = farr[j];
  ++colnum_ttype;
  
  //read vcirc
  fits_get_colnum(fptr,CASEINSEN,ttype[colnum_ttype],&colnum,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' column # = %d\n",ttype[colnum_ttype],colnum);
#endif
  fits_read_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,&nulval,farr,&anynul,&status);
  if(status)
    fits_report_error(stderr,status);
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    fprintf(stderr,"'%s' anynul = %d\n",ttype[colnum_ttype],anynul);
#endif
  for(j=0;j<HaloChunkSize;++j)
    chunkHalos[j].vcirc = farr[j];
  ++colnum_ttype;
  
  //read px,py,pz
  for(i=0;i<3;++i)
    {
      fits_get_colnum(fptr,CASEINSEN,ttype[colnum_ttype],&colnum,&status);
      if(status)
	fits_report_error(stderr,status);
#ifdef DEBUG
      if(DEBUG_LEVEL > 1)
	fprintf(stderr,"'%s' column # = %d\n",ttype[colnum_ttype],colnum);
#endif
      fits_read_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,&nulval,farr,&anynul,&status);
      if(status)
	fits_report_error(stderr,status);
#ifdef DEBUG
      if(DEBUG_LEVEL > 1)
	fprintf(stderr,"'%s' anynul = %d\n",ttype[colnum_ttype],anynul);
#endif
      for(j=0;j<HaloChunkSize;++j)
	chunkHalos[j].pos[i] = farr[j];
      ++colnum_ttype;
    }
      
  //read vx, vy, vz
  for(i=0;i<3;++i)
    {
      fits_get_colnum(fptr,CASEINSEN,ttype[colnum_ttype],&colnum,&status);
      if(status)
	fits_report_error(stderr,status);
#ifdef DEBUG
      if(DEBUG_LEVEL > 1)
	fprintf(stderr,"'%s' column # = %d\n",ttype[colnum_ttype],colnum);
#endif
      fits_read_col(fptr,TFLOAT,colnum,firstrow,firstelem,nelements,&nulval,farr,&anynul,&status);
      if(status)
	fits_report_error(stderr,status);
#ifdef DEBUG
      if(DEBUG_LEVEL > 1)
	fprintf(stderr,"'%s' anynul = %d\n",ttype[colnum_ttype],anynul);
#endif
      for(j=0;j<HaloChunkSize;++j)
	chunkHalos[j].vel[i] = farr[j];
      ++colnum_ttype;
    }
  
  fits_close_file(fptr,&status);
  if(status)
    fits_report_error(stderr,status);
  
  free(buff);
      
#ifdef DEBUG
  if(DEBUG_LEVEL > 2)
    for(i=0;i<13;++i)
      fprintf(stderr,"%d: index = %ld, radius,mass = %f|%e, px,py,pz = %f|%f|%f, vx,vy,vz = %f|%f|%f\n",
	      ThisTask,chunkHalos[i].index,chunkHalos[i].radius,chunkHalos[i].mass,chunkHalos[i].pos[0],chunkHalos[i].pos[1],chunkHalos[i].pos[2],
	      chunkHalos[i].vel[0],chunkHalos[i].vel[1],chunkHalos[i].vel[2]);
#endif

}
